{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Training a SSD MobileNet V2 for Pedestrian Detection with TLT\n",
    "\n",
    "This notebook shows how to train a SSD MobileNet V2 object detector for pedestrian detection with [Nvidia's Transfer Learning Toolkit](https://developer.nvidia.com/transfer-learning-toolkit).\n",
    "TLT is a useful tool for train classification and object detection models with various pretrained backbone. The trained models with TLT can be easily integrate and deploy to DeepStream.\n",
    "Note that this notebook and specs files are just a modification of examples that exist on the [TLT docker container](https://ngc.nvidia.com/catalog/containers/nvidia:tlt-streamanalytics).\n",
    "\n",
    "0. [Set up env variables](#head-0)\n",
    "1. [Prepare dataset and pre-trained model](#head-1) <br>\n",
    "    1.1 [Prepare tfrecords from kitti format dataset](#head-1-1) <br>\n",
    "    1.2 [Download pre-trained model](#head-1-2) <br>\n",
    "2. [Provide training specification](#head-2)\n",
    "3. [Run TLT training](#head-3)\n",
    "4. [Evaluate trained models](#head-4)\n",
    "5. [Prune trained models](#head-5)\n",
    "6. [Retrain pruned models](#head-6)\n",
    "7. [Evaluate retrained model](#head-7)\n",
    "8. [Visualize inferences](#head-8)\n",
    "9. [Deploy](#head-9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables <a class=\"anchor\" id=\"head-0\"></a>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "print(\"Please replace the variable with your key.\")\n",
    "%set_env KEY=[YOUR NGC API KEY]\n",
    "%set_env USER_EXPERIMENT_DIR=/experiment_dir/ped_ssd_mobilenet_v2\n",
    "%set_env DATASET_DIR=/experiment_dir/dataset\n",
    "%set_env SPECS_DIR=/repo/training/tlt/pedestrian_detection/specs\n",
    "%set_env TF_FORCE_GPU_ALLOW_GROWTH=true\n",
    "!mkdir -p $USER_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using [Oxford Town Center Dataset](https://www.robots.ox.ac.uk/ActiveVision/Research/Projects/2009bbenfold_headpose/project.html) for training our pedestrian detection. First we will download the video file and annotatoin `csv` file with `download_towncenter_video_and_labels.sh` script and then create images and kitti format annotation files with [ffmpeg](https://ffmpeg.org/) and `extract_kitti_labels.py` scripts. note that TLT will not perform dynamic resizing in training time and every image should be resize to a certain size offline (for SSD detectord the width and height of images should be a multiply of 32). Also for training object detectors the annotations should be in KITTI format. for further information about creating KITTI format annotations checkout [TLT getting started guide](https://docs.nvidia.com/metropolis/TLT/tlt-getting-started-guide/index.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download Oxford Town Center Dataset video file and csv annotation file\n",
    "!bash ./download_towncenter_video_and_labels.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract images from video, resize them and save in images directory with `ffmpeg`\n",
    "!mkdir $DATASET_DIR/images\n",
    "!ffmpeg -i $DATASET_DIR/TownCentreXVID.avi -q:v 1 -start_number 0 -vf scale=320:320 -frames:v 4501 $DATASET_DIR/images/%d.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create KITTI annotation file and save in labels directory\n",
    "!python3 extract_kitti_labels.py --annotation_path $DATASET_DIR/TownCentre-groundtruth.top --image_width 320 --image_height 320"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "!ls -lh $DATASET_DIR/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Prepare tfrecords from kitti format dataset <a class=\"anchor\" id=\"head-1-1\"></a>\n",
    "\n",
    "* Update the tfrecords spec file to take in your kitti format dataset\n",
    "* Create the tfrecords using the tlt-dataset-convert \n",
    "* TFRecords only need to be generated once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TFrecords conversion spec file for training\")\n",
    "!cat $SPECS_DIR/ssd_tfrecords_towncenter_trainval.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a new directory for the output tfrecords dump.\n",
    "# !mkdir -p $USER_EXPERIMENT_DIR/tfrecords\n",
    "#KITTI trainval\n",
    "!tlt-dataset-convert -d $SPECS_DIR/ssd_tfrecords_towncenter_trainval.txt \\\n",
    "                     -o $DATA_DOWNLOAD_DIR/tfrecords/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $DATASET_DIR/tfrecords/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Download pre-trained model <a class=\"anchor\" id=\"head-1-2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tlt_pretrained_object_detection:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $USER_EXPERIMENT_DIR/pretrained_mobilenet_v2/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tlt_pretrained_object_detection:mobilenet_v2 --dest $USER_EXPERIMENT_DIR/pretrained_mobilenet_v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -lh $USER_EXPERIMENT_DIR/pretrained_mobilenet_v2/tlt_pretrained_object_detection_vmobilenet_v2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Provide training specification <a class=\"anchor\" id=\"head-2\"></a>\n",
    "* Tfrecords for the train datasets\n",
    "    * In order to use the newly generated tfrecords, update the dataset_config parameter in the spec file at `$SPECS_DIR/ssd_train_resnet18_kitti.txt` \n",
    "    * Update the fold number to use for evaluation. In case of random data split, please use fold 0 only\n",
    "    * For sequence wise you may use any fold generated from the dataset convert tool\n",
    "* Pre-trained models\n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $SPECS_DIR/ped_ssd_mobilenet_v2_train.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Run TLT training <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_unpruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"To run with multigpu, please change --gpus based on the number of available GPUs in your machine.\")\n",
    "!tlt-train ssd -e $SPECS_DIR/ped_ssd_mobilenet_v2_train.txt \\\n",
    "               -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \\\n",
    "               -k $KEY \\\n",
    "               -m $USER_EXPERIMENT_DIR/pretrained_mobilenet_v2/tlt_pretrained_object_detection_vmobilenet_v2/mobilenet_v2.hdf5 \\\n",
    "               --gpus 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "# Note csv epoch number is 1 less than model file epoch. For example, epoch 79 in csv corresponds to _080.tlt\n",
    "!cat $USER_EXPERIMENT_DIR/experiment_dir_unpruned/ssd_training_log_mobilenet_v2.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the epoch number which has best performance\n",
    "%set_env EPOCH=060"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Evaluate trained models <a class=\"anchor\" id=\"head-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tlt-evaluate ssd -e $SPECS_DIR/ped_ssd_mobilenet_v2_train.txt \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt \\\n",
    "                  -k $KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Prune trained models <a class=\"anchor\" id=\"head-5\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion (`Only for resnets as they have element wise operations or MobileNets.`)\n",
    "* Threshold for pruning.\n",
    "* A key to save and load the model\n",
    "* Output directory to store the model\n",
    "\n",
    "Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold value depends on the dataset and the model. `0.5` in the block below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tlt-prune -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt \\\n",
    "           -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/ssd_mobilenet_v2_pruned.tlt \\\n",
    "           -eq intersection \\\n",
    "           -pth 0.5 \\\n",
    "           -k $KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $USER_EXPERIMENT_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Retrain pruned models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Printing the retrain spec file. \n",
    "# Here we have updated the spec file to include the newly pruned model as a pretrained weights.\n",
    "!cat $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $USER_EXPERIMENT_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tlt-train ssd --gpus 1 \\\n",
    "               -e $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt \\\n",
    "               -r $USER_EXPERIMENT_DIR/experiment_dir_retrain \\\n",
    "               -m $USER_EXPERIMENT_DIR/experiment_dir_pruned/ssd_mobilenet_v2_pruned.tlt \\\n",
    "               -k $KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlht $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "# Note csv epoch number is 1 less than model file epoch. For example, epoch 79 in csv corresponds to _080.tlt\n",
    "!cat $USER_EXPERIMENT_DIR/experiment_dir_retrain/ssd_training_log_mobilenet_v2.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the epoch number which has best performance\n",
    "%set_env EPOCH=035"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Evaluate retrained model <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tlt-evaluate ssd -e $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt \\\n",
    "                  -k $KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Visualize inferences <a class=\"anchor\" id=\"head-8\"></a>\n",
    "In this section, we run the tlt-infer tool to generate inferences on the trained models and visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tlt-infer ssd -i $DATASET_DIR/images \\\n",
    "               -o $USER_EXPERIMENT_DIR/ssd_infer_images \\\n",
    "               -e $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt \\\n",
    "               -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt \\\n",
    "               -l $USER_EXPERIMENT_DIR/ssd_infer_labels \\\n",
    "               -k $KEY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `tlt-infer` tool produces two outputs. \n",
    "1. Overlain images in `$USER_EXPERIMENT_DIR/ssd_infer_images`\n",
    "2. Frame by frame bbox labels in kitti format located in `$USER_EXPERIMENT_DIR/ssd_infer_labels`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['USER_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx / num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'ssd_infer_images' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Deploy! <a class=\"anchor\" id=\"head-9\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!mkdir -p $USER_EXPERIMENT_DIR/export\n",
    "# Export in FP32 mode. Change --data_type to fp16 for FP16 mode\n",
    "!tlt-export ssd -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt \\\n",
    "                -k $KEY \\\n",
    "                -o $USER_EXPERIMENT_DIR/export/ped_ssd_mobilenet_v2_epoch_$EPOCH.etlt \\\n",
    "                -e $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt \\\n",
    "                --batch_size 1 \\\n",
    "                --data_type fp16\n",
    "\n",
    "# Uncomment to export in INT8 mode (generate calibration cache file). \\\n",
    "# !tlt-export ssd -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/ssd_mobilenet_v2_epoch_$EPOCH.tlt  \\\n",
    "#                 -o $USER_EXPERIMENT_DIR/export/ped_ssd_mobilenet_v2_epoch_$EPOCH.etlt \\\n",
    "#                 -e $SPECS_DIR/ped_ssd_mobilenet_v2_retrain.txt \\\n",
    "#                 -k $KEY \\\n",
    "#                 --cal_image_dir  $USER_EXPERIMENT_DIR/data/testing/image_2 \\\n",
    "#                 --data_type int8 \\\n",
    "#                 --batch_size 1 \\\n",
    "#                 --batches 10 \\\n",
    "#                 --cal_cache_file $USER_EXPERIMENT_DIR/export/cal.bin  \\\n",
    "#                 --cal_data_file $USER_EXPERIMENT_DIR/export/cal.tensorfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Note:` In this example, for ease of execution we restrict the number of calibrating batches to 10. TLT recommends the use of at least 10% of the training dataset for int8 calibration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### You can use this `.etlt` file directly in DeepStream or first create a `Tensor RT Engine` and then use the `engine` in DeepStream"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
